{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import os\n",
    "%matplotlib inline\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "part_dict = ['nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear',\\\n",
    "             'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow',\\\n",
    "             'left_wrist', 'right_wrist', 'left_hip', 'right_hip', 'left_knee',\\\n",
    "             'right_knee', 'left_ankle', 'right_ankle']\n",
    "\n",
    "COCO_dir = '/group/wanbo/HOI_plus/data/coco/images/train2014'\n",
    "\n",
    "\n",
    "def showKps(kps):\n",
    "    #kps: 51, show coco keypoints\n",
    "    #kps = kps.reshape((17,3))\n",
    "    v = kps[:,2]\n",
    "    skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13],[6, 12], [7, 13], [6, 7], [6, 8],\\\n",
    "                [7, 9], [8, 10], [9, 11], [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]]\n",
    "\n",
    "    for sk in skeleton:\n",
    "        if v[sk[0]-1]>0 and v[sk[1]-1]>0:\n",
    "            plt.plot([kps[sk[0]-1,0], kps[sk[1]-1,0]], [kps[sk[0]-1,1], kps[sk[1]-1,1]], 'r-')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'pickle' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-8e9722afa697>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0mhoi_blob_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhoi_blob_path\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m'saved_hoi_blobs3.pkl'\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m \u001b[0mhoi_data\u001b[0m\u001b[0;34m=\u001b[0m \u001b[0mpickle\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhoi_blob_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'rb'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      7\u001b[0m \u001b[0mhoi_data\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhoi_data\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'pickle' is not defined"
     ]
    }
   ],
   "source": [
    "import matplotlib.patches as patches\n",
    "\n",
    "hoi_blob_path = '/root/backup/Research/HOI_plus/Outputs/fine_net3/debug/saved_hoi_blobs/'\n",
    "\n",
    "hoi_blob_path = hoi_blob_path + 'saved_hoi_blobs3.pkl'\n",
    "hoi_data= pickle.load(open(hoi_blob_path, 'rb'))\n",
    "hoi_data = hoi_data[0]\n",
    "\n",
    "im_id = hoi_data['id']\n",
    "im_path = os.path.join(COCO_dir, 'COCO_train2014_'+str(im_id).zfill(12)+'.jpg')\n",
    "\n",
    "print(os.path.isdir(COCO_dir))\n",
    "#img = Image.open(im_path)\n",
    "img = cv2.imread(im_path)\n",
    "img = img[:,:,[2,1,0]]\n",
    "plt.figure(figsize=(8,8))\n",
    "if hoi_data['flipped']:\n",
    "    img = img[:,::-1]\n",
    "plt.imshow(img)\n",
    "\n",
    "print (hoi_data.keys())\n",
    "# for k, v in hoi_data.items():\n",
    "#     if type(v)==np.ndarray:\n",
    "#         print(k, v.shape)\n",
    "\n",
    "#for i in range(union_boxes.shape[0]):\n",
    "human_boxes = hoi_data['human_boxes'][hoi_data['interaction_human_inds']]\n",
    "object_boxes = hoi_data['object_boxes'][hoi_data['interaction_object_inds']]\n",
    "\n",
    "print('human_boxe.shape ', human_boxes.shape )\n",
    "print('obj_boxes.shape ', object_boxes.shape)\n",
    "\n",
    "union_boxes = hoi_data['union_boxes']\n",
    "gt_union_heatmap = hoi_data['gt_union_heatmap']\n",
    "union_mask = hoi_data['union_mask']\n",
    "\n",
    "#print('union mask.shape:', union_mask.shape)\n",
    "#print(human_boxes)\n",
    "#print(object_boxes)\n",
    "\n",
    "for i in range(union_boxes.shape[0]):\n",
    "    plt.figure(figsize=(16,8))\n",
    "    plt.subplot(241)\n",
    "    plt.imshow(gt_union_heatmap[i,:17].sum(axis=0))\n",
    "    plt.title('pose heatmap')\n",
    "    \n",
    "    pts = np.zeros((17,3))\n",
    "    for j in range(17):\n",
    "        y, x = np.unravel_index(np.argmax(gt_union_heatmap[i,j]), (64,64))\n",
    "        pts[j,0] = x\n",
    "        pts[j,1] = y\n",
    "        pts[j,2] = np.max(gt_union_heatmap[i,j])\n",
    "    \n",
    "    plt.subplot(242)\n",
    "    plt.imshow(gt_union_heatmap[i,17])\n",
    "    plt.title('human proposal')\n",
    "    \n",
    "    plt.subplot(243)\n",
    "    plt.imshow(gt_union_heatmap[i,18])\n",
    "    plt.title('object proposal')\n",
    "    \n",
    "    x1,y1,x2,y2 = (union_boxes[i,1:]/hoi_data['image_info'][2]).astype(np.int32)\n",
    "    \n",
    "    hx1, hy1, hx2, hy2 = (human_boxes[i, 1:]/hoi_data['image_info'][2]).astype(np.int32)\n",
    "    ox1, oy1, ox2, oy2 = (object_boxes[i, 1:]/hoi_data['image_info'][2]).astype(np.int32)\n",
    "    \n",
    "    scale_x = 64*1. / (x2-x1)\n",
    "    scale_y = 64*1. / (y2-y1)\n",
    "    #print (x1,y1,x2,y2)\n",
    "    \n",
    "    hx1 = ((hx1-x1)*scale_x).astype(np.int32)\n",
    "    hx2 = ((hx2-x1)*scale_x).astype(np.int32)\n",
    "    hy1 = ((hy1-y1)*scale_y).astype(np.int32)\n",
    "    hy2 = ((hy2-y1)*scale_y).astype(np.int32)\n",
    "    \n",
    "    #print( hx1, hy1, hx2, hy2)\n",
    "    ox1 = ((ox1-x1)*scale_x).astype(np.int32)\n",
    "    ox2 = ((ox2-x1)*scale_x).astype(np.int32)\n",
    "    oy1 = ((oy1-y1)*scale_y).astype(np.int32)\n",
    "    oy2 = ((oy2-y1)*scale_y).astype(np.int32)\n",
    "    #print(ox1, oy1, ox2, oy2)\n",
    "    \n",
    "    unimg = img[y1:y2, x1:x2, :]\n",
    "    unimg = cv2.resize(unimg, (64,64))\n",
    "    ax = plt.subplot(244)\n",
    "    plt.imshow(unimg)\n",
    "    showKps(pts)\n",
    "    #print(pts[:,2])\n",
    "    if np.max(pts[:,2])<1:\n",
    "        plt.title('predicted pose')\n",
    "    else:\n",
    "        plt.title('gt pose')\n",
    "    \n",
    "    rect = patches.Rectangle((hx1, hy1), (hx2-hx1), (hy2-hy1),linewidth=3,edgecolor='r',facecolor='none')\n",
    "    ax.add_patch(rect)\n",
    "    \n",
    "    rect = patches.Rectangle((ox1, oy1), (ox2-ox1), (oy2-oy1),linewidth=3,edgecolor='r',facecolor='none')\n",
    "    ax.add_patch(rect)\n",
    "    #plt.plot(pts[:,0], pts[:,1], 'ro')\n",
    "    \n",
    "    plt.subplot(245)\n",
    "    plt.imshow(union_mask[i,0])\n",
    "    plt.title('human proposal')\n",
    "    \n",
    "    plt.subplot(246)\n",
    "    plt.imshow(union_mask[i,4])\n",
    "    plt.title('union mask residual {}'.format(np.min(union_mask[i,4])))\n",
    "    \n",
    "    plt.subplot(247)\n",
    "    plt.imshow(union_mask[i,2])\n",
    "    plt.title('union mask ori')\n",
    "    \n",
    "    plt.subplot(248)\n",
    "    plt.imshow(union_mask[i,3])\n",
    "    plt.title('union mask noise ')\n",
    "    \n",
    "print(hoi_data['flipped'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(21, 37)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(640, 427)\n",
      "325 92 474 322\n"
     ]
    }
   ],
   "source": [
    "print(img.size)\n",
    "print(x1, y1, x2, y2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(16, 5)\n"
     ]
    }
   ],
   "source": [
    "print(union_boxes.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1069.       ,  800.       ,    1.6701461], dtype=float32)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hoi_data['image_info']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(16, 5)\n",
      "(15, 5)\n"
     ]
    }
   ],
   "source": [
    "print(hoi_data['union_boxes'].shape)\n",
    "print(hoi_data['object_boxes'].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1068.8"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1.67*640"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.8735363"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
